package com.example.demo.advisor;

import org.springframework.ai.chat.client.ChatClientMessageAggregator;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import org.springframework.ai.chat.prompt.Prompt;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;

/**
 * 自定义 Re2 Advisor
 * 可提高大型语言模型的推理能力
 */
public class ReReadingAdvisor implements CallAdvisor, StreamAdvisor {

    /**
     * 执行请求前，改写 Prompt
     *
     * @param advisedRequest
     * @return
     */
    private ChatClientRequest before(ChatClientRequest advisedRequest) {
        String userPrompt = advisedRequest.prompt().getUserMessage().getText();

        advisedRequest.context().put("re2_input_query", userPrompt);

        String text = """
                %s
                Read the question again: %s
                """.formatted(userPrompt, userPrompt);

        Prompt newPrompt = advisedRequest.prompt().augmentUserMessage(userPrompt);
        return new ChatClientRequest(newPrompt, advisedRequest.context());
    }

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest advisedRequest, CallAdvisorChain callAdvisorChain) {
        return callAdvisorChain.nextCall(this.before(advisedRequest));
    }

    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest advisedRequest, StreamAdvisorChain streamAdvisorChain) {
        advisedRequest = before(advisedRequest);
        Flux<ChatClientResponse> chatClientResponseFlux = streamAdvisorChain.nextStream(advisedRequest);
        return chatClientResponseFlux;
    }

    @Override
    public int getOrder() {
        return 0;
    }

    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }
}
